package study.datastructure.tree;

import lombok.val;

import java.util.*;
import java.util.stream.Collectors;

public class AVL<T extends Comparable<T>> {

    public class Node {
        int key;
        T val;
        Node left;
        Node right;
        // 以该节点为根的树的高度
        int height;
        int balanceFactor;

        public Node(T val) {
            this.val = val;
            this.height = 1; // 默认高度为1
        }

        public Node(int key, T val) {
            this.key = key;
            this.val = val;
        }

        @Override
        public String toString() {
            return String.format("val:%s, height:%d, balanceFactor:%d", this.val, this.height, this.balanceFactor);
        }
    }

    private Node root;
    private Integer size;

    public AVL() {
        this.root = null;
        this.size = 0;
    }

    public boolean isEmpty() {
        return this.size == 0;
    }

    public Integer getSize() {
        return this.size;
    }

    // 验证是否是二叉搜索树
    public boolean isBinaryTree() {
        List<T> list = new ArrayList<>();
        midTraversal(list, this.root);
        for (int i = 1; i < list.size(); i++) {
            if (list.get(i - 1).compareTo(list.get(i)) > 0) {
                return false;
            }
        }
        return true;
    }

    // 获取当前节点的高度
    public int getHeight(Node node) {
        return node == null ? 0 : node.height;
    }

    // 获取当前节点的平衡因子
    public int getBalanceFactor(Node node) {
        return node == null ? 0 : getHeight(node.left) - getHeight(node.right);
    }

    public boolean isBalanceTree() {
        return isBalanceTree(this.root);
    }

    // 判断以node为节点的树是否是平衡二叉树
    private boolean isBalanceTree(Node node) {
        if (node == null) {
            return true;
        }
        int balanceFactor = Math.abs(getBalanceFactor(node));

        if (balanceFactor > 1) {
            return false;
        } else {
            return isBalanceTree(node.left) && isBalanceTree(node.right);
        }
    }


    /**
     * 左旋转
     *
     * @param red 当前节点
     * @return 旋转后的根节点
     */
    public Node leftRotate(Node red) {
        Node yellow = red.right;

        red.right = yellow.left;
        yellow.left = red;
        // 更新节点的高度
        updateHeight(red); //先更新 更新后较低节点的高度
        updateHeight(yellow);//再更新 更新后较高节点的高度
        return yellow;
    }

    /**
     * 右旋转
     *
     * @param red 当前节点
     * @return 旋转后的根节点
     */
    public Node rightRotate(Node red) {
        Node yellow = red.left;
        red.left = yellow.right;    //换爹

        yellow.right = red;  //上位
        // 更新节点的高度
        updateHeight(red);//先更新 更新后较低节点的高度
        updateHeight(yellow);//再更新 更新后较高节点的高度
        return yellow;
    }

    /**
     * 左右旋转
     *
     * @param node 当前节点
     * @return 旋转后的根节点
     */
    public Node leftRightRotate(Node node) {
        leftRotate(node.left);
        return rightRotate(node);
    }

    /**
     * 右左旋转
     *
     * @param node 当前节点
     * @return 旋转后的根节点
     */
    public Node rightLeftRotate(Node node) {
        rightRotate(node.right);
        rightRotate(node.right);
        return leftRightRotate(node);
    }

    private void updateHeight(Node x) {
        x.height = Math.max(getHeight(x.left), getHeight(x.right)) + 1;
    }


    // 添加
    public void add(T val) {
        this.root = add(this.root, val);
    }

    private Node add(Node node, T val) {
        if (node == null) {
            this.size++;
            Node leafNode = new Node(val);
            return leafNode;
        }
        // 当前节点的值小于添加的val, 因此val做右孩子
        if (node.val.compareTo(val) < 0) {
            node.right = add(node.right, val);
            // 更新高度
            updateHeight(node);
        } else if (node.val.compareTo(val) > 0) {
            node.left = add(node.left, val);
            // 更新高度
            updateHeight(node);
        } else {
            return node;
        }
        return  balance(node);
    }

    public void put(int key, T val) {
        this.root = doPut(this.root, key, val);
    }

    private Node doPut(Node node, int key, T val) {
        //找到空位
        if (node == null) {
            return new Node(key,val);
        }

        //key 已存在 更新
        if (key == node.key) {
            node.val = val;
            return node;
        }

        if (key < node.key) {
            node.left=doPut(node.left, key, val);
        } else {
            node.right=doPut(node.right, key, val);
        }
        updateHeight(node);
        return balance(node);

    }

    /**
     * 自平衡 旋转
     *
     * @param node
     * @return 返回旋转后根节点
     */
    private Node balance(Node node) {
       Node resNode =null;
        int balanceFactor = getBalanceFactor(node);
        if (balanceFactor > 1 && getBalanceFactor(node.left) >= 0) {//LL
            // 右
            resNode = rightRotate(node);
        } else if (balanceFactor > 1 && getBalanceFactor(node.left) < 0) {//LR
            // 左右
            leftRightRotate(node);
        } else if (balanceFactor < -1 && getBalanceFactor(node.right) <= 0) {//RR
            // 左
            resNode = leftRotate(node);
        } else if (balanceFactor < -1 && getBalanceFactor(node.right) > 0) {//RL
            // 右左
            rightLeftRotate(node);
        }
        return resNode;
    }

    // 删除操作
    // 删除树中的val
    public void remove(T val) {
        if (!contains(val)) {
            return;
        }
        this.root = remove(this.root, val);
    }
    public void remove(int key) {
        this.root = doRemove(this.root, key);
    }

    private Node doRemove(Node node, int key) {
        if (node == null) {
            return null;
        }
        //key 已存在 更新
        if (key == node.key) {

            return node;
        }

        return null;
    }

    /**
     * 删除val
     *
     * @param node
     * @param val
     * @return
     */
    public Node remove(Node node, T val) {
        // 递归终止条件
        if (node == null) {
            return null;
        }
        Node resNode = null;
        if (node.val.compareTo(val) == 0) {
            this.size--;
            if (node.right == null) {
                // 右子树为空
                resNode = node.left;
            } else if (node.left == null) {
                // 左子树为空
                resNode = node.right;
            } else {
                // 左右子树都不为空
                // 1. 找到删除节点的后继
                Node suffixNode = getMin(node.right);
                // 2. 删除后继
                suffixNode.right = remove(node.right, suffixNode.val);
                // 3. 连接
                suffixNode.left = node.left;
                // 返回删除后的根
                resNode = suffixNode;
            }
        } else if (node.val.compareTo(val) < 0) {
            node.right = remove(node.right, val);
            resNode = node;
        } else {
            node.left = remove(node.left, val);
            resNode = node;
        }

        // 删除节点可能为叶子结点
        if (resNode == null) {
            return null;
        }
        // 更新高度
        updateHeight(resNode);
        int balanceFactor = getBalanceFactor(resNode);

        return balance(node);
    }

    private Node getMin(Node node) {
        // 递归终止条件
        if (node.left == null) {
            return node;
        }
        // 递归操作
        return getMin(node.left);
    }

    public String midTraversal() {
        List<T> res = new ArrayList<>();
        midTraversal(res, this.root);
        return res.stream().map(item -> item.toString()).collect(Collectors.joining(","));
    }

    /**
     * 中序遍历
     *
     * @param result
     * @param node   当前节点
     * @return
     */
    private void midTraversal(List<T> result, Node node) {
        if (node == null) {
            return;
        }
        midTraversal(result, node.left);
        result.add(node.val);
        midTraversal(result, node.right);
    }

    public void showTree() {
        showTree(this.root);
    }

    private void showTree(Node node) {
        if (node == null) {
            return;
        }
        showTree(node.left);
        System.out.println(node);
        showTree(node.right);
    }

    // 查询是否存在val
    public boolean contains(T val) {
        return contains(this.root, val);
    }

    private boolean contains(Node node, T val) {
        // 递归的终止条件
        // 查询到低也没有找到
        if (node == null) {
            return false;
        }
        // 递归操作
        if (node.val.compareTo(val) == 0) {
            return true;
        } else if (node.val.compareTo(val) < 0) {
            return contains(node.right, val);
        } else {
            return contains(node.left, val);
        }
    }

    // 测试
    public static void main(String[] args) {
        AVL<Integer> avlTree = new AVL<>();

        // 测试添加节点操作
        avlTree.add(10);
        avlTree.add(20);
        avlTree.add(30);
        avlTree.add(40);
        avlTree.add(50);
        avlTree.add(25);

        // 输出树的中序遍历结果，验证二叉搜索树的性质
        System.out.println("中序遍历结果: " + avlTree.midTraversal());

        // 测试是否为平衡二叉树
        System.out.println("是否为平衡二叉树: " + avlTree.isBalanceTree());

        // 测试删除节点操作
        avlTree.remove(30);

        // 再次输出中序遍历结果
        System.out.println("删除节点后中序遍历结果: " + avlTree.midTraversal());

        // 再次测试是否为平衡二叉树
        System.out.println("删除节点后是否为平衡二叉树: " + avlTree.isBalanceTree());

        // 测试查询节点是否存在
        System.out.println("是否存在25节点: " + avlTree.contains(25));
        System.out.println("是否存在30节点: " + avlTree.contains(30));

        // 输出树的节点信息（通过showTree方法，这里只是简单展示节点的基本属性情况）
        System.out.println("树的节点信息展示:");
        avlTree.showTree();
    }
}
